import numpy as np

def predict_action(state: np.ndarray):
    """
    策略2：布林带 + KDJ + 波动率 + 成交量 组合策略
    输入:
        state: np.ndarray of shape (b, 32, 20)
    输出:
        action: np.ndarray of shape (b, 1)，取值为 [-1, 0, 1]
    """
    b, l, c = state.shape
    assert l == 32 and c == 20

    # 指标索引
    IDX = {
        "close": 4,
        "volume": 5,
        "kdj_k": 13,
        "atr": 14,
        "bb_upper": 15
    }

    close = state[:, -1, IDX["close"]]
    bb_upper = state[:, -1, IDX["bb_upper"]]
    kdj_k = state[:, -1, IDX["kdj_k"]]
    kdj_k_prev = state[:, -2, IDX["kdj_k"]]
    atr = state[:, -1, IDX["atr"]]
    atr_prev = state[:, -2, IDX["atr"]]
    volume = state[:, -1, IDX["volume"]]
    volume_prev = state[:, -2, IDX["volume"]]

    # 多头信号
    long_cond = (
        (close < bb_upper * 0.98) &
        (kdj_k < 30) & (kdj_k > kdj_k_prev) &
        (atr > atr_prev) &
        (volume > volume_prev)
    )

    # 空头信号
    short_cond = (
        (close > bb_upper * 1.02) &
        (kdj_k > 70) & (kdj_k < kdj_k_prev) &
        (atr > atr_prev) &
        (volume > volume_prev)
    )

    action = np.zeros(b, dtype=np.int32)
    action[long_cond] = 1
    action[short_cond] = -1

    return action.reshape(-1, 1)
